# @Time: 2025/4/22 21:45
# @Author: Yuanqing He
# @Email: heyuanqing007@163.com
import os
import zipfile

import requests
import json
import time

NODE_IPs = ["223.2.39.61"]  # 替换为实际的IP地址
# NODE_IPs = ["127.0.0.1"]  # 替换为实际的IP地址
MODELS = []
Current_IP = None


def register(node_ips):
    """
    注册节点并获取所有节点上的模型列表

    :param node_ips: 节点的IP地址列表
    :return: 包含每个节点IP地址及其模型列表的字典
    """
    global MODELS
    all_model_lists = {}

    for ip in node_ips:
        url = f"http://{ip}:8000/get-model-list"
        try:
            response = requests.get(url)
            response.raise_for_status()  # 如果响应状态码不是200，会抛出异常
            model_list = response.json()
            all_model_lists[ip] = model_list
            print(f"成功获取节点 {ip} 的模型列表")
        except requests.exceptions.RequestException as e:
            print(f"无法获取节点 {ip} 的模型列表: {e}")
    MODELS = all_model_lists


def initialize_model(model_name):
    """
    根据模型名称初始化模型

    :param model_name: 要初始化的模型名称
    :return: 初始化结果
    """
    global MODELS
    global Current_IP
    for ip, models in MODELS.items():
        for model in models:
            if model.get("Model_Name") == model_name:
                init_url = f"http://{ip}:8000/initialize-model/{model_name}"
                try:
                    response = requests.get(init_url)
                    response.raise_for_status()  # 如果响应状态码不是200，会抛出异常
                    print(f"成功初始化模型 {model_name} 在节点 {ip}")
                    Current_IP = ip
                    return response.json()
                except requests.exceptions.RequestException as e:
                    print(f"无法初始化模型 {model_name} 在节点 {ip}: {e}")
                    return None
    print(f"未找到模型 {model_name}")
    return None


def upload_file(url, file_path):
    """
    上传文件到指定的URL

    :param url: 上传文件的URL
    :param file_path: 要上传的文件路径
    :return: 响应的文本内容
    """
    with open(file_path, 'rb') as file:
        files = {'file': file}
        response = requests.post(url, files=files)
        if response.status_code == 200:
            print(f"文件 {file_path} 上传成功")
        else:
            print(f"文件 {file_path} 上传失败")
    return response.text


def upload_model_files(file_list):
    """
    根据当前模型所在的节点上传文件

    :param file_list: 文件列表，每个文件信息包括文件名和路径
    :return: 上传结果
    """
    global Current_IP
    if Current_IP is None:
        print("当前没有初始化模型，无法上传文件")
        return

    for file_info in file_list:
        file_name = file_info.get("file_name")
        file_path = file_info.get("file_path")
        if not file_name or not file_path:
            print(f"文件信息不完整: {file_info}")
            continue

        upload_url = f"http://{Current_IP}:8000/upload-file/{file_name}"
        upload_file(upload_url, file_path)


def run_model(params):
    """
    调用当前激活模型

    :param params: 调用模型的参数
    :return: 模型调用结果
    """
    global Current_IP
    if Current_IP is None:
        print("当前没有初始化模型，无法调用模型")
        return None

    run_url = f"http://{Current_IP}:8000/run-model"
    try:
        response = requests.post(run_url, json=params)
        response.raise_for_status()  # 如果响应状态码不是200，会抛出异常
        print(f"成功调用模型在节点 {Current_IP}")
        return response.json()
    except requests.exceptions.RequestException as e:
        print(f"无法调用模型在节点 {Current_IP}: {e}")
        return None


import urllib
import re


def get_filename_from_cd(content_disposition):
    """
    从 Content-Disposition 头中提取文件名，支持普通和 RFC5987 编码格式
    """
    # 尝试匹配编码格式（如 UTF-8''filename.ext）
    utf8_filename_regex = re.compile(r"filename\*=UTF-8''([\w%\-\.]+)", re.IGNORECASE)
    match = utf8_filename_regex.search(content_disposition)
    if match:
        return urllib.parse.unquote(match.group(1))

    # 尝试匹配普通格式（如 filename="example.txt"）
    ascii_filename_regex = re.compile(r'filename="?([^"]+)"?', re.IGNORECASE)
    match = ascii_filename_regex.search(content_disposition)
    if match:
        return match.group(1)

    return None


def download_results(output_file_path=None):
    """
    下载模型结果

    :param output_file_path: 下载文件的本地保存路径（可选）
    :return: 下载结果
    """
    global Current_IP
    if Current_IP is None:
        print("当前没有初始化模型，无法下载结果")
        return

    download_url = f"http://{Current_IP}:8000/download-results"
    try:
        response = requests.get(download_url, stream=True)
        response.raise_for_status()  # 如果响应状态码不是200，会抛出异常

        # 从响应头中提取文件名
        content_disposition = response.headers.get('Content-Disposition')
        file_name = get_filename_from_cd(content_disposition)
        if not file_name:
            file_name = "model_results.zip"

        if output_file_path:
            # 如果用户提供了保存路径，则使用用户提供的路径
            file_path = output_file_path
        else:
            # 如果用户没有提供保存路径，则使用当前文件夹和提取的文件名
            file_path = os.path.join(os.getcwd(), file_name)

        with open(file_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)

        if file_path.lower().endswith('.zip'):
            extract_dir = os.path.splitext(file_path)[0]  # 使用ZIP文件名创建解压目录
            os.makedirs(extract_dir, exist_ok=True)
            with zipfile.ZipFile(file_path, 'r') as zip_ref:
                zip_ref.extractall(extract_dir)
            os.remove(file_path)

        print(f"数据下载成功")
    except requests.exceptions.RequestException as e:
        print(f"无法下载模型结果: {e}")



def print_hi(name):
    # Use a breakpoint in the code line below to debug your script.
    print(f'Hi, {name}')  # Press Ctrl+F8 to toggle the breakpoint.


# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    # 用户输入节点的IP地址列表

    # 获取所有节点上的模型列表
    # 1.注册模型节点并获取模型列表
    register(NODE_IPs)
    # 2.初始化模型
    initialize_model('Multiply_numbers')
    # 3.配置输入数据
    file_list = [
        {"file_name": "num1.txt", "file_path": "num1.txt"},
        {"file_name": "num3.txt", "file_path": "num3.txt"}
    ]
    upload_model_files(file_list)
    # 4.调用模型
    params = {
        "number1": "num1.txt",  # 替换为你的文件路径
        "number2": 4
    }
    result = run_model(params)
    if result:
        print(f"模型调用结果: {result}")

    # 5.下载模型结果
    download_results()
    print_hi('PyCharm')

# See PyCharm help at https://www.jetbrains.com/help/pycharm/
